#include <numeric>
#include <chrono>
#include "evdeploy/runtime/ascend/atlas_common.h"
#include "evdeploy/utils/report/tool_report.h"

namespace ev
{

    AscendModel::AscendModel()
    {
        EVLOG(INFO) << "this is a common atals engine";
    }

    AscendModel::~AscendModel()
    {
        UnInit();
    }

    EVStatus AscendModel::UnInit()
    {
        if (m_init_flag)
        {
            EVLOG(INFO) << "in uninit func";
            TRY_RELEASE_POINTER(m_input_dataset, aclmdlDestroyDataset);
            for (int i = 0; i < m_input_sizes.size(); ++i)
            {
                TRY_RELEASE_POINTER(m_input_data_buffers[i], aclDestroyDataBuffer);
                TRY_RELEASE_POINTER(m_input_buffers[i], aclrtFree);
            }

            TRY_RELEASE_POINTER(m_output_dataset, aclmdlDestroyDataset);
            for (int i = 0; i < m_output_buffers_host.size(); i++)
            {
                if (m_output_buffers_host[i] != nullptr)
                {
                    delete m_output_buffers_host[i];
                }
            }
            aclmdlDestroyDesc(m_model_desc);
            if (m_model_id > 0)
            {
                aclmdlUnload(m_model_id);
            }
            aclrtDestroyStream(m_acl_stream);
            aclrtDestroyContext(m_acl_context);
            // m_chw_wrappers.clear();
            m_init_flag = false;
            return EV_SUCCESS;
        }
        return EV_REPEATED_OP;
    }

    // init model and infos
    EVStatus AscendModel::Init(const char *net, void **config)
    {
        if (m_init_flag)
        {
            EVLOG(WARNING) << "call init mutiple times";
            return EV_REPEATED_OP;
        }

        EVLOG(INFO) << "in init func";
        ToolReport report;
        // report info
        int code = report.report(3,
                                 "ev_model_common-arm-Ascend",
                                 "V1.0.0",
                                 "common model interface, altas version");
        if (code == 0)
        {
            EVLOG(INFO) << "report done successully";
        }

        SetContext(config[0]);
        auto ret = aclrtSetCurrentContext(*m_acl_context);
        ret = aclrtCreateStream(&m_acl_stream);

        if (ret != 0)
        {
            EVLOG(ERROR) << "failed to create  stream";
        }

        ret = aclrtGetRunMode(&m_acl_run_mode);
        if (ret != 0)
        {
            EVLOG(ERROR) << "failed to get acl run mode";
        }

        LoadModel(net);

        m_model_desc = aclmdlCreateDesc();
        ret = aclmdlGetDesc(m_model_desc, m_model_id);

        // �~N��~O~V��~S�~E�信�~A���~L�~H~F�~E~M��~S�~E���~D��~P
        m_input_dataset = aclmdlCreateDataset();
        int input_num = aclmdlGetNumInputs(m_model_desc);
        EVLOG(INFO) << "Model input num is : " << input_num;
        for (int i = 0; i < input_num; ++i)
        {
            aclmdlIODims input_dim;
            aclmdlGetInputDims(m_model_desc, i, &input_dim);
            m_input_dims.push_back(input_dim);
            std::string dims_info = "Model input " + std::to_string(i) + " has shape : (";
            for (int k = 0; k < input_dim.dimCount; ++k)
            {
                dims_info += std::to_string(input_dim.dims[k]) + ",";
            }
            dims_info[dims_info.size() - 1] = ')';
            dims_info = dims_info + " and data size " + std::to_string(aclmdlGetInputSizeByIndex(m_model_desc, i)) + " Bytes";
            EVLOG(INFO) << dims_info;
            void *input_buffer = nullptr;
            auto ret = aclrtMalloc(&input_buffer, aclmdlGetInputSizeByIndex(m_model_desc, i), ACL_MEM_MALLOC_NORMAL_ONLY);
            auto input_data_buffer = aclCreateDataBuffer(input_buffer, aclmdlGetInputSizeByIndex(m_model_desc, i));
            aclmdlAddDatasetBuffer(m_input_dataset, input_data_buffer);
            m_input_buffers.push_back(input_buffer);
            m_input_data_buffers.push_back(input_data_buffer);
            m_input_sizes.push_back(aclmdlGetInputSizeByIndex(m_model_desc, i));
        }
        // �~N��~O~V��~S�~G�信�~A���~L�~H~F�~E~M��~S�~G���~D��~P
        m_output_dataset = aclmdlCreateDataset();
        int output_num = aclmdlGetNumOutputs(m_model_desc);
        EVLOG(INFO) << "Model output num is : " << output_num;
        for (int i = 0; i < output_num; ++i)
        {
            aclmdlIODims output_dim;
            aclmdlGetOutputDims(m_model_desc, i, &output_dim);
            m_output_dims.push_back(output_dim);
            std::string dims_info = "Model output " + std::to_string(i) + " has shape : (";
            for (int k = 0; k < output_dim.dimCount; ++k)
            {
                dims_info += std::to_string(output_dim.dims[k]) + ",";
            }
            dims_info[dims_info.size() - 1] = ')';
            dims_info = dims_info + " and data size " + std::to_string(aclmdlGetOutputSizeByIndex(m_model_desc, i)) + " Bytes";
            EVLOG(INFO) << dims_info;

            void *output_buffer = nullptr;
            auto ret = aclrtMalloc(&output_buffer, aclmdlGetOutputSizeByIndex(m_model_desc, i), ACL_MEM_MALLOC_NORMAL_ONLY);
            auto output_data_buffer = aclCreateDataBuffer(output_buffer, aclmdlGetOutputSizeByIndex(m_model_desc, i));
            aclmdlAddDatasetBuffer(m_output_dataset, output_data_buffer);
            m_output_buffers.push_back(output_buffer);
            m_output_data_buffers.push_back(output_data_buffer);
            m_output_sizes.push_back(aclmdlGetOutputSizeByIndex(m_model_desc, i));
            if (m_acl_run_mode == ACL_HOST)
            {
                EVLOG(INFO) << "atlas run in host mode";
                void *buffer = new char[aclmdlGetOutputSizeByIndex(m_model_desc, i)];
                m_output_buffers_host.push_back(buffer);
            }
        }
        // �~H~[建��~S�~E��~U��~M���~S��~X
        m_init_flag = true;
        return EV_SUCCESS;
    }

    // run inference
    EVStatus AscendModel::Infer(EVModelData *in, EVModelData *out)
    {
        REC_TIME(t0);
        aclrtSetCurrentContext(*m_acl_context);
        if (m_input_dims.size() != in->mat_num)
        {
            EVLOG(ERROR) << "Model requires: " << m_input_dims.size() << " inputs but " << in->mat_num << " is provided.";
            return EV_FAIL;
        }

        for (int i = 0; i < in->mat_num; ++i)
        {
            if (in->mat[i].data_size != m_input_sizes[i])
            {
                EVLOG(ERROR) << i << " size mismatch " << in->mat[i].data_size << " vs " << m_input_sizes[i];
                return EV_FAIL;
            }

            if (in->mat[i].loc == EV_DATA_HOST)
            {
                aclrtMemcpy(m_input_buffers[i], m_input_sizes[i], in->mat[i].data, m_input_sizes[i], ACL_MEMCPY_HOST_TO_DEVICE);
            }
            else
            {
                aclrtMemcpy(m_input_buffers[i], m_input_sizes[i], in->mat[i].data, m_input_sizes[i], ACL_MEMCPY_DEVICE_TO_DEVICE);
            }
        }

        auto ret = aclmdlExecute(m_model_id, m_input_dataset, m_output_dataset);
        out->mat_num = m_output_dims.size();
        out->mat = new EVMatData[m_output_dims.size()];
        for (int i = 0; i < m_output_dims.size(); ++i)
        {
            out->mat[i].loc = EV_DATA_HOST;
            out->mat[i].data_size = m_output_sizes[i];
            aclmdlIODims output_dim = m_output_dims[i];
            for (int j = 0; j < output_dim.dimCount; ++j)
            {
                out->mat[i].dims.push_back(output_dim.dims[j]);
            }
            out->mat[i].desc = "";
            if (m_acl_run_mode == ACL_DEVICE)
            {
                out->mat[i].data = m_output_buffers[i];
            }
            else
            {
                aclrtMemcpy(m_output_buffers_host[i], m_output_sizes[i], m_output_buffers[i], m_output_sizes[i], ACL_MEMCPY_DEVICE_TO_HOST);
                out->mat[i].data = m_output_buffers_host[i];
            }
        }

        REC_TIME(t1);
        EVLOG(INFO) << "inference time(ms):" << RUN_TIME(t1 - t0);

        return EV_SUCCESS;
    }

    EVStatus AscendModel::OptimizeModel(const char *model_name)
    {
        return EV_SUCCESS;
    }

    EVStatus AscendModel::SetIO(const std::vector<std::string> &inputs, const std::vector<std::string> &outputs)
    {
        return EV_SUCCESS;
    }

    CREATE_INFER_INS(AscendModel)
    DESTROY_INFER_INS(AscendModel)

} // namespace ev
